# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/26 14:47 
@Author  : ZhangShenao 
@File    : 3.ChatModel流式输出.py
@Desc    : ChatModel流式输出
"""
from datetime import datetime

from langchain_community.chat_models.moonshot import MoonshotChat
from langchain_core.prompts import ChatPromptTemplate

# 编排Prompt
chat_prompt_template = ChatPromptTemplate.from_messages([
    ('system', '你是Moonshot开发的AI聊天机器人，请根据用户的提问，生成相应的回答。现在的时间是：{now}'),
    ('human', '{query}'),
]).partial(now=datetime.now())

# 生成消息
msg = chat_prompt_template.invoke({
    'query': '将帮我生成一首宋词',
})

# 创建ChatModel大语言模型
llm = MoonshotChat(model='moonshot-v1-8k')

# stream()函数,获取流式输出
message_chunks = llm.stream(msg)

# 打印流式输出内容
for chunk in message_chunks:
    print(chunk.content, flush=True, end="")
